package com.evaluation_system.Util;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.GeneralSecurityException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;

/**
 * PDF 文件加密解密工具类 (AES-256-CBC 模式)
 */
public class PdfEncryptor {
    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/CBC/PKCS5Padding";
    private static final int KEY_SIZE = 256;
    private static final int IV_SIZE = 16; // AES块大小固定为16字节

    // ------------------- 加密相关 -------------------
    /**
     * 生成 AES 密钥 (256位)
     */
    public static SecretKey generateKey() throws NoSuchAlgorithmException {
        KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
        keyGen.init(KEY_SIZE);
        return keyGen.generateKey();
    }

    /**
     * 生成随机初始化向量 (IV)
     */
    public static IvParameterSpec generateIv() {
        byte[] iv = new byte[IV_SIZE];
        new SecureRandom().nextBytes(iv);
        return new IvParameterSpec(iv);
    }

    /**
     * 加密 PDF 字节数据
     * @param pdfBytes 原始PDF字节数组
     * @param key      AES密钥
     * @param iv       初始化向量
     * @return 加密后的字节数组
     */
    public static byte[] encrypt(byte[] pdfBytes, SecretKey key, IvParameterSpec iv)
            throws GeneralSecurityException {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, key, iv);
        return cipher.doFinal(pdfBytes);
    }

    // ------------------- 解密相关 -------------------
    /**
     * 解密 PDF 字节数据
     * @param encryptedBytes 加密后的字节数组
     * @param key           AES密钥 (必须与加密时相同)
     * @param iv            初始化向量 (必须与加密时相同)
     * @return 解密后的原始PDF字节数组
     */
    public static byte[] decrypt(byte[] encryptedBytes, SecretKey key, IvParameterSpec iv)
            throws GeneralSecurityException {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, key, iv);
        return cipher.doFinal(encryptedBytes);
    }

    // ------------------- 密钥转换工具 -------------------
    /**
     * 将密钥转换为Base64字符串 (便于存储或传输)
     */
    public static String keyToString(SecretKey key) {
        return Base64.getEncoder().encodeToString(key.getEncoded());
    }

    /**
     * 从Base64字符串恢复密钥
     */
    public static SecretKey stringToKey(String keyStr) {
        byte[] decodedKey = Base64.getDecoder().decode(keyStr);
        return new SecretKeySpec(decodedKey, 0, decodedKey.length, ALGORITHM);
    }

    /**
     * 将IV转换为Base64字符串
     */
    public static String ivToString(IvParameterSpec iv) {
        return Base64.getEncoder().encodeToString(iv.getIV());
    }

    /**
     * 从Base64字符串恢复IV
     */
    public static IvParameterSpec stringToIv(String ivStr) {
        byte[] ivBytes = Base64.getDecoder().decode(ivStr);
        return new IvParameterSpec(ivBytes);
    }
}